package scheduler

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"gitee.com/war2001/NiuMaEasyGo/global"
	"gitee.com/war2001/NiuMaEasyGo/nm_core/queue"
	"github.com/google/uuid"
	"github.com/robfig/cron/v3"
	"github.com/streadway/amqp"
	"go.uber.org/zap"
)

// DistributedTask 分布式任务
type DistributedTask struct {
	ID          string                 `json:"id"`           // 任务ID
	Name        string                 `json:"name"`         // 任务名称
	Topic       string                 `json:"topic"`        // 任务主题
	Params      map[string]interface{} `json:"params"`       // 任务参数
	Timeout     int                    `json:"timeout"`      // 任务超时时间(秒)
	RetryCount  int                    `json:"retry_count"`  // 重试次数
	RetryDelay  int                    `json:"retry_delay"`  // 重试延迟(秒)
	ScheduledAt time.Time              `json:"scheduled_at"` // 计划执行时间
	CreatedAt   time.Time              `json:"created_at"`   // 创建时间
	Status      string                 `json:"status"`       // 任务状态: pending, running, success, failed
}

// DistributedScheduler 分布式调度器
type DistributedScheduler struct {
	cron             *cron.Cron
	tasks            map[string]cron.EntryID
	mutex            sync.RWMutex
	ctx              context.Context
	cancel           context.CancelFunc
	wg               sync.WaitGroup
	timezone         *time.Location
	running          bool
	exchangeName     string
	routingKeyPrefix string
	mq               *queue.RabbitMQ
	handlers         map[string]TaskHandler
	isPublisher      bool
	isConsumer       bool
}

// TaskHandler 任务处理函数类型
type TaskHandler func(ctx context.Context, task *DistributedTask) error

// DistributedConfig 分布式调度器配置
type DistributedConfig struct {
	Timezone         string         // 时区
	Location         *time.Location // 时区位置
	ExchangeName     string         // 交换机名称
	RoutingKeyPrefix string         // 路由键前缀
	IsPublisher      bool           // 是否为发布者
	IsConsumer       bool           // 是否为消费者
}

// NewDistributedScheduler 创建分布式调度器
func NewDistributedScheduler(config *DistributedConfig) (*DistributedScheduler, error) {
	// 设置时区
	location := time.Local
	if config.Location != nil {
		location = config.Location
	} else if config.Timezone != "" {
		var err error
		location, err = time.LoadLocation(config.Timezone)
		if err != nil {
			return nil, err
		}
	}

	// 创建上下文
	ctx, cancel := context.WithCancel(context.Background())

	// 创建分布式调度器
	scheduler := &DistributedScheduler{
		cron:             cron.New(cron.WithLocation(location), cron.WithSeconds()),
		tasks:            make(map[string]cron.EntryID),
		ctx:              ctx,
		cancel:           cancel,
		timezone:         location,
		running:          false,
		exchangeName:     config.ExchangeName,
		routingKeyPrefix: config.RoutingKeyPrefix,
		handlers:         make(map[string]TaskHandler),
		isPublisher:      config.IsPublisher,
		isConsumer:       config.IsConsumer,
	}

	return scheduler, nil
}

// Start 启动分布式调度器
func (s *DistributedScheduler) Start() error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.running {
		return nil
	}

	// 检查消息队列是否已初始化
	if mq, ok := global.MessageQueue.(*queue.RabbitMQ); ok {
		s.mq = mq
	} else {
		return fmt.Errorf("消息队列未初始化或类型不匹配")
	}

	// 如果是发布者，启动cron调度器
	if s.isPublisher {
		s.cron.Start()
	}

	// 如果是消费者，启动任务消费
	if s.isConsumer {
		if err := s.startConsuming(); err != nil {
			return err
		}
	}

	s.running = true
	zap.L().Info("分布式调度器已启动",
		zap.Bool("publisher", s.isPublisher),
		zap.Bool("consumer", s.isConsumer))

	return nil
}

// Stop 停止分布式调度器
func (s *DistributedScheduler) Stop() {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if !s.running {
		return
	}

	// 停止cron调度器
	if s.isPublisher && s.cron != nil {
		s.cron.Stop()
	}

	// 取消上下文
	s.cancel()

	// 等待所有任务完成
	s.wg.Wait()

	s.running = false
	zap.L().Info("分布式调度器已停止")
}

// AddTask 添加定时任务
func (s *DistributedScheduler) AddTask(name, topic, spec string, params map[string]interface{}) error {
	if !s.isPublisher {
		return fmt.Errorf("当前实例不是发布者，无法添加定时任务")
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	// 检查任务是否已存在
	for taskName, _ := range s.tasks {
		if taskName == name {
			return fmt.Errorf("任务 %s 已存在", name)
		}
	}

	// 包装任务执行函数
	wrapper := func() {
		task := &DistributedTask{
			ID:          uuid.New().String(),
			Name:        name,
			Topic:       topic,
			Params:      params,
			Timeout:     300, // 默认5分钟超时
			RetryCount:  3,   // 默认重试3次
			RetryDelay:  60,  // 默认重试延迟1分钟
			ScheduledAt: time.Now(),
			CreatedAt:   time.Now(),
			Status:      "pending",
		}

		// 发布任务到消息队列
		if err := s.publishTask(task); err != nil {
			zap.L().Error("发布任务失败",
				zap.String("task", name),
				zap.String("topic", topic),
				zap.Error(err))
		}
	}

	// 添加到cron
	entryID, err := s.cron.AddFunc(spec, wrapper)
	if err != nil {
		return err
	}

	// 记录任务
	s.tasks[name] = entryID
	zap.L().Info("任务已添加到分布式调度器",
		zap.String("task", name),
		zap.String("topic", topic),
		zap.String("spec", spec))

	return nil
}

// RegisterHandler 注册任务处理函数
func (s *DistributedScheduler) RegisterHandler(topic string, handler TaskHandler) error {
	if !s.isConsumer {
		return fmt.Errorf("当前实例不是消费者，无法注册处理函数")
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	// 检查处理函数是否已存在
	if _, exists := s.handlers[topic]; exists {
		return fmt.Errorf("主题 %s 的处理函数已存在", topic)
	}

	// 注册处理函数
	s.handlers[topic] = handler
	zap.L().Info("任务处理函数已注册", zap.String("topic", topic))

	return nil
}

// PublishTask 手动发布任务
func (s *DistributedScheduler) PublishTask(name, topic string, params map[string]interface{}) error {
	task := &DistributedTask{
		ID:          uuid.New().String(),
		Name:        name,
		Topic:       topic,
		Params:      params,
		Timeout:     300, // 默认5分钟超时
		RetryCount:  3,   // 默认重试3次
		RetryDelay:  60,  // 默认重试延迟1分钟
		ScheduledAt: time.Now(),
		CreatedAt:   time.Now(),
		Status:      "pending",
	}

	return s.publishTask(task)
}

// publishTask 发布任务到消息队列
func (s *DistributedScheduler) publishTask(task *DistributedTask) error {
	// 序列化任务
	data, err := json.Marshal(task)
	if err != nil {
		return fmt.Errorf("序列化任务失败: %v", err)
	}

	// 构建路由键
	routingKey := s.routingKeyPrefix + "." + task.Topic

	// 发布消息
	err = s.mq.PublishWithRoutingKey(routingKey, data)
	if err != nil {
		return fmt.Errorf("发布任务到消息队列失败: %v", err)
	}

	zap.L().Info("任务已发布",
		zap.String("id", task.ID),
		zap.String("name", task.Name),
		zap.String("topic", task.Topic),
		zap.String("routing_key", routingKey))

	return nil
}

// startConsuming 开始消费任务
func (s *DistributedScheduler) startConsuming() error {
	// 为每个主题创建队列和绑定
	for topic := range s.handlers {
		queueName := fmt.Sprintf("task_queue_%s", topic)
		routingKey := s.routingKeyPrefix + "." + topic

		// 声明队列
		if err := s.declareQueue(queueName, routingKey); err != nil {
			return err
		}

		// 注册消费者
		if err := s.registerConsumer(queueName, topic); err != nil {
			return err
		}
	}

	return nil
}

// declareQueue 声明队列并绑定到交换机
func (s *DistributedScheduler) declareQueue(queueName, routingKey string) error {
	// 获取通道
	ch, err := s.mq.GetChannel()
	if err != nil {
		return fmt.Errorf("获取通道失败: %v", err)
	}

	// 声明队列
	_, err = ch.QueueDeclare(
		queueName, // 队列名称
		true,      // 持久化
		false,     // 自动删除
		false,     // 排他性
		false,     // 非等待
		nil,       // 参数
	)
	if err != nil {
		return fmt.Errorf("声明队列失败: %v", err)
	}

	// 绑定队列到交换机
	err = ch.QueueBind(
		queueName,      // 队列名称
		routingKey,     // 路由键
		s.exchangeName, // 交换机名称
		false,          // 非等待
		nil,            // 参数
	)
	if err != nil {
		return fmt.Errorf("绑定队列到交换机失败: %v", err)
	}

	return nil
}

// registerConsumer 注册消费者
func (s *DistributedScheduler) registerConsumer(queueName, topic string) error {
	handler := s.handlers[topic]

	// 创建消费者
	consumerTag := fmt.Sprintf("task_consumer_%s_%s", topic, uuid.New().String())

	// 注册消费者
	err := s.mq.Consume(consumerTag, false, func(delivery amqp.Delivery) error {
		// 解析任务
		var task DistributedTask
		if err := json.Unmarshal(delivery.Body, &task); err != nil {
			zap.L().Error("解析任务失败",
				zap.Error(err),
				zap.ByteString("body", delivery.Body))
			return err
		}

		// 更新任务状态
		task.Status = "running"

		// 创建上下文
		ctx, cancel := context.WithTimeout(s.ctx, time.Duration(task.Timeout)*time.Second)
		defer cancel()

		// 执行任务
		zap.L().Info("开始执行任务",
			zap.String("id", task.ID),
			zap.String("name", task.Name),
			zap.String("topic", task.Topic))

		s.wg.Add(1)
		defer s.wg.Done()

		err := handler(ctx, &task)
		if err != nil {
			// 任务执行失败
			task.Status = "failed"
			zap.L().Error("任务执行失败",
				zap.String("id", task.ID),
				zap.String("name", task.Name),
				zap.String("topic", task.Topic),
				zap.Error(err))

			// 检查是否需要重试
			if task.RetryCount > 0 {
				task.RetryCount--
				task.Status = "pending"

				// 重新发布任务
				go func() {
					time.Sleep(time.Duration(task.RetryDelay) * time.Second)
					if err := s.publishTask(&task); err != nil {
						zap.L().Error("重新发布任务失败",
							zap.String("id", task.ID),
							zap.String("name", task.Name),
							zap.Error(err))
					}
				}()
			}

			return err
		}

		// 任务执行成功
		task.Status = "success"
		zap.L().Info("任务执行成功",
			zap.String("id", task.ID),
			zap.String("name", task.Name),
			zap.String("topic", task.Topic))

		return nil
	})

	if err != nil {
		return fmt.Errorf("注册消费者失败: %v", err)
	}

	zap.L().Info("任务消费者已注册",
		zap.String("queue", queueName),
		zap.String("topic", topic),
		zap.String("consumer", consumerTag))

	return nil
}

// RemoveTask 移除定时任务
func (s *DistributedScheduler) RemoveTask(name string) {
	if !s.isPublisher {
		zap.L().Warn("当前实例不是发布者，无法移除定时任务")
		return
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	if entryID, exists := s.tasks[name]; exists {
		s.cron.Remove(entryID)
		delete(s.tasks, name)
		zap.L().Info("任务已从分布式调度器移除", zap.String("task", name))
	}
}

// ListTasks 列出所有定时任务
func (s *DistributedScheduler) ListTasks() []string {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	tasks := make([]string, 0, len(s.tasks))
	for name := range s.tasks {
		tasks = append(tasks, name)
	}
	return tasks
}

// GetTaskCount 获取任务数量
func (s *DistributedScheduler) GetTaskCount() int {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return len(s.tasks)
}

// IsRunning 检查调度器是否运行中
func (s *DistributedScheduler) IsRunning() bool {
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	return s.running
}

// GetTaskEntry 获取任务条目
func (s *DistributedScheduler) GetTaskEntry(name string) *cron.Entry {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	if entryID, exists := s.tasks[name]; exists {
		entry := s.cron.Entry(entryID)
		return &entry
	}
	return nil
}
